//
//  Data+AYExtension.swift
//  Adorbee
//
//  Created by amovision on 2024/11/27.
//
import CommonCrypto
import Compression
import UIKit

public extension Data {
  
  enum Algorithm {
    case lz4
    case zlib
    case lzma
    case lzfse

    var compressionAlgorithm: compression_algorithm {
      switch self {
      case .lz4: return COMPRESSION_LZ4
      case .zlib: return COMPRESSION_ZLIB
      case .lzma: return COMPRESSION_LZMA
      case .lzfse: return COMPRESSION_LZFSE
      }
    }
  }
  
  func toInt<T: FixedWidthInteger>(at index: Int) -> T {
    guard index + MemoryLayout<T>.size <= count else { return 0 }
    return withUnsafeBytes {
      $0.load(fromByteOffset: index, as: T.self).bigEndian
    }
  }
    
  var int8: Int8 { return toInt(at: 0) }
  var uint8: UInt8 { return toInt(at: 0) }
  var int16: Int16 { return toInt(at: 0) }
  var uint16: UInt16 { return toInt(at: 0) }
  var int32: Int32 { return toInt(at: 0) }
  var uint32: UInt32 { return toInt(at: 0) }
  var int64: Int64 { return toInt(at: 0) }
  var uint64: UInt64 { return toInt(at: 0) }
  var ns: NSData { self as NSData }
  init<T: FixedWidthInteger>(_ int: T) {
    // 获取整数的字节数
    let byteCount = MemoryLayout<T>.size
    self = Data(count: byteCount)
    withUnsafeMutableBytes { bytes in
      bytes.storeBytes(of: int, as: T.self)
    }
  }
    
  func string(encoding: String.Encoding) -> String? {
    return String(data: self, encoding: encoding)
  }
    
  init?(hex: String) {
    // 去除空格和换行符
    let cleanedHex = hex.replacingOccurrences(of: " ", with: "").replacingOccurrences(of: "\n", with: "")
    
    // 确保字符串长度是偶数
    guard cleanedHex.count % 2 == 0 else {
      return nil
    }
    var data = Data(capacity: cleanedHex.count / 2)
    var index = cleanedHex.startIndex
    for _ in 0..<cleanedHex.count / 2 {
      let nextIndex = cleanedHex.index(index, offsetBy: 2)
      let byteString = String(cleanedHex[index..<nextIndex])
      if let byte = UInt8(byteString, radix: 16) {
        data.append(byte)
      } else {
        return nil // 如果转换失败，返回nil
      }
      index = nextIndex
    }
    self = data
  }

  var hexString: String {
    return map { String(format: "%02x", $0) }.joined()
  }
  
  var md5String: String {
    var digest = [UInt8](repeating: 0, count: Int(CC_MD5_DIGEST_LENGTH))
    _ = withUnsafeBytes { bytes in
      CC_MD5(bytes.baseAddress, CC_LONG(self.count), &digest)
    }
    return digest.map { String(format: "%02hhx", $0) }.joined()
  }
  
  var sha256String: String {
    var digest = [UInt8](repeating: 0, count: Int(CC_SHA256_DIGEST_LENGTH))
    _ = withUnsafeBytes { bytes in
      CC_SHA256(bytes.baseAddress, CC_LONG(self.count), &digest)
    }
    return digest.map { String(format: "%02hhx", $0) }.joined()
  }

  func hmacSHA256(key: Data) -> String {
    var hmac = [UInt8](repeating: 0, count: Int(CC_SHA256_DIGEST_LENGTH))
    key.withUnsafeBytes { keyBytes in
      self.withUnsafeBytes { dataBytes in
        CCHmac(CCHmacAlgorithm(kCCHmacAlgSHA256), keyBytes.baseAddress, key.count, dataBytes.baseAddress, self.count, &hmac)
      }
    }
    return hmac.map { String(format: "%02x", $0) }.joined()
  }
    
  func encryptWithRSA(publicKeyBase64: String) -> Data? {
    guard let publicKeyData = Data(base64Encoded: publicKeyBase64),
          let publicKeyStr = String(data: publicKeyData, encoding: .utf8)
    else {
      print("Invalid Base64 string")
      return nil
    }
    var publicKeyString = publicKeyStr
    publicKeyString = publicKeyString.replacingOccurrences(of: "-----BEGIN RSA PUBLIC KEY-----", with: "")
    publicKeyString = publicKeyString.replacingOccurrences(of: "-----END RSA PUBLIC KEY-----", with: "")
    publicKeyString = publicKeyString.replacingOccurrences(of: "\r", with: "")
    publicKeyString = publicKeyString.replacingOccurrences(of: "\n", with: "")
    publicKeyString = publicKeyString.replacingOccurrences(of: " ", with: "")
    guard let result = Data(base64Encoded: publicKeyString) else {
      print("Invalid Base64 string")
      return nil
    }
    guard let publicKey = SecKeyCreateWithData(result as CFData,
                                               [kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
                                                kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
                                                kSecAttrKeySizeInBits as String: 2048] as CFDictionary,
                                               nil)
    else {
      print("Error creating public key")
      return nil
    }
        
    var error: Unmanaged<CFError>?
    guard let encryptedData = SecKeyCreateEncryptedData(publicKey,
                                                        SecKeyAlgorithm.rsaEncryptionOAEPSHA256,
                                                        self as CFData,
                                                        &error)
    else {
      print("Error encrypting data: \(error!.takeRetainedValue() as Error)")
      return nil
    }
        
    return encryptedData as Data
  }
    
  func aes256Decrypt(key: Data, iv: Data) -> Data? {
    // 检查密钥长度
    guard key.count == 16 || key.count == 24 || key.count == 32 else {
      return nil
    }
    // 检查 IV 长度
    guard iv.count == 16 || iv.isEmpty else {
      return nil
    }
            
    var result: Data?
    let bufferSize = count + kCCBlockSizeAES128
    var buffer = Data(count: bufferSize)
    var encryptedSize: size_t = 0
            
    // 执行解密
    let cryptStatus = buffer.withUnsafeMutableBytes { bufferBytes in
      self.withUnsafeBytes { dataBytes in
        key.withUnsafeBytes { keyBytes in
          iv.withUnsafeBytes { ivBytes in
            CCCrypt(CCOperation(kCCDecrypt),
                    CCAlgorithm(kCCAlgorithmAES),
                    CCOptions(kCCOptionPKCS7Padding),
                    keyBytes.baseAddress,
                    key.count,
                    ivBytes.baseAddress,
                    dataBytes.baseAddress,
                    self.count,
                    bufferBytes.baseAddress,
                    bufferSize,
                    &encryptedSize)
          }
        }
      }
    }
            
    // 检查解密状态
    if cryptStatus == kCCSuccess {
      buffer.removeSubrange(encryptedSize..<buffer.count) // 移除多余的字节
      result = buffer
    }
            
    return result
  }
  
  func compress(algorithm: Algorithm) -> Data {
    return withUnsafeBytes { (sourcePtr: UnsafeRawBufferPointer) -> Data in
      let sourceSize = count
      let destinationSize = sourceSize
      let destinationBuffer = UnsafeMutablePointer<UInt8>.allocate(capacity: destinationSize)
      defer { destinationBuffer.deallocate() }

      let compressedSize = compression_encode_buffer(
        destinationBuffer,
        destinationSize,
        sourcePtr.baseAddress!.assumingMemoryBound(to: UInt8.self),
        sourceSize,
        nil,
        algorithm.compressionAlgorithm
      )

      return Data(bytes: destinationBuffer, count: compressedSize)
    }
  }

  func decompress(algorithm: Algorithm) -> Data {
    return withUnsafeBytes { (sourcePtr: UnsafeRawBufferPointer) -> Data in
      let sourceSize = count
      let destinationSize = sourceSize * 10 // Assume a maximum compression ratio of 10
      let destinationBuffer = UnsafeMutablePointer<UInt8>.allocate(capacity: destinationSize)
      defer { destinationBuffer.deallocate() }

      let decompressedSize = compression_decode_buffer(
        destinationBuffer,
        destinationSize,
        sourcePtr.baseAddress!.assumingMemoryBound(to: UInt8.self),
        sourceSize,
        nil,
        algorithm.compressionAlgorithm
      )

      return Data(bytes: destinationBuffer, count: decompressedSize)
    }
  }
}
